package org.csploit.android.net.metasploit;

import org.csploit.android.R;
import org.csploit.android.core.Logger;
import org.csploit.android.core.System;
import org.csploit.android.net.Target;
import org.csploit.android.net.reference.CVE;
import org.csploit.android.net.reference.OSVDB;
import org.csploit.android.net.reference.Reference;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;

/**
 * this class store MSF Exploit
 */
public class MsfExploit extends Target.Exploit
{
  private Ranking rank = Ranking.Normal;
  private ArrayList<String> targets = new ArrayList<String>();
  private Collection<Author> authors = new LinkedList<Author>();
  private Collection<String> platforms = new LinkedList<String>();
  private Collection<String> architectures = new LinkedList<String>();
  private String license = null;


  // RPC fields
  private boolean synced = false;
  private String currentTarget = null;
  private Payload currentPayload = null;
  private int version = 0;
  private Collection<String> jobs = new LinkedList<String>();
  private Collection<Option> options = new ArrayList<Option>();
  private ArrayList<Payload> payloads = new ArrayList<Payload>();

  public enum Ranking {
    Manual(R.color.msf_rank_manual, R.string.msf_rank_manual),
    Low(R.color.msf_rank_low, R.string.msf_rank_low),
    Average(R.color.msf_rank_average, R.string.msf_rank_average),
    Normal(R.color.msf_rank_normal, R.string.msf_rank_normal),
    Good(R.color.msf_rank_good, R.string.msf_rank_good),
    Great(R.color.msf_rank_great, R.string.msf_rank_great),
    Excellent(R.color.msf_rank_excellent, R.string.msf_rank_excellent);

    final int color;
    final int stringID;

    Ranking(int color, int stringID) {
      this.color = color;
      this.stringID = stringID;
    }

    public int getColor() {
      return color;
    }

    public int getStringID() {
      return stringID;
    }

    public static Ranking valueOf(int value) {
      switch (value) {
        case 0:
          return Manual;
        case 100:
          return Low;
        case 200:
          return Average;
        default:
        case 300:
          return Normal;
        case 400:
          return Good;
        case 500:
          return Great;
        case 600:
          return Excellent;
      }
    }
  }

  public MsfExploit(String name, Target.Port port) {
    super(name, "http://www.rapid7.com/db/modules/" + name);
    this.id = name;
    this.port = port;
    refresh();
  }

  public MsfExploit(String name, String summary, String description, Ranking rank,
                    ArrayList<String> targets, Collection<Author> authors, Collection<String> platforms,
                    Collection<String> architectures, Collection<Reference> references) {
    this.id = name;
    this.name = name;
    this.url = "http://www.rapid7.com/db/modules/" + name;
    this.summary = summary;
    this.description = description;
    this.rank = rank;
    this.targets = targets;
    this.authors = authors;
    this.platforms = platforms;
    this.architectures = architectures;
    if(references != null) {
      this.references.addAll(references);
    }

    refresh();
  }

  public synchronized void refresh() {
    if(synced || System.getMsfRpc() == null)
      return;
    try {
      retrieveInfos();
      retrieveOptions();
      retrievePayloads();
      synced = true;
    } catch (IOException e) {
      System.errorLogging(e);
    } catch (RPCClient.MSFException e) {
      Logger.error(name +": MSFException: "+e.getMessage());
    }
  }

  public boolean isSynced() {
    return synced;
  }

  @SuppressWarnings("unchecked")
  private void retrieveOptions() throws IOException, RPCClient.MSFException {
    Object res = System.getMsfRpc().call("module.options", "exploit", name);

    if(res == null)
      return;

    for(Map.Entry<String, HashMap<String, Object>> entry : ((HashMap<String,HashMap<String, Object>>)res).entrySet()) {
      String name = entry.getKey();
      Option option;

      try {
        option = new Option(name, entry.getValue());
      } catch (IllegalArgumentException e) {
        System.errorLogging(e);
        continue;
      }

      if(name.equals("RHOST")) {
        if(parent != null)
          option.setValue(parent.getAddress().getHostAddress());
      } else if(name.equals("RPORT")) {
        if(port != null)
          option.setValue(Integer.toString(port.getNumber()));
      }

      options.add(option);
    }
  }

  @SuppressWarnings("unchecked")
  private void retrieveInfos() throws IOException, RPCClient.MSFException {
    HashMap<String,Object> map = (HashMap<String,Object>) System.getMsfRpc().call("module.info", "exploit", name);
    Object res;

    // use the name as summary, it should be quite short
    summary = (String)map.get("name");
    // remove all newlines and tab chars, android UI will show it in the best way possible.
    description = ((String)map.get("description")).replace("\n", "").replace("\t", "");
    // get rank
    rank = Ranking.valueOf((Integer)map.get("rank"));
    // get version
    try {
      version = Integer.parseInt((String)map.get("version"));
    } catch (NumberFormatException e) {
      Logger.warning(description + ": unknown version");
      version = 0;
    }

    // get references
    res = map.get("references");
    if(res != null) {
      for (ArrayList<String> entry : ((ArrayList<ArrayList<String>>) res)) {
        String type = entry.get(0);
        String id = entry.get(1);

        if (type.equals("CVE")) {
          references.add(new CVE(id));
        } else if (type.equals("OSVDB")) {
          references.add(new OSVDB(Integer.parseInt(id)));
        }
      }
    }

    // get targets
    res = map.get("targets");
    if(res != null) {
      targets.clear();
      for(Map.Entry<Integer, String> e : ((Map<Integer,String>)res).entrySet()) {
        targets.add(e.getValue());
      }
      
      Integer def = (Integer) map.get("default_target");
      if(def != null) {
        currentTarget = targets.get(def);
      } else if(!targets.isEmpty()) {
        currentTarget = targets.get(0);
      } else {
        currentTarget = null;
      }
    }

    // get authors
    res = map.get("authors");
    if(res != null) {
      authors.clear();

      for(String nameAndEmail : (ArrayList<String>)res) {
        authors.add(Author.fromString(nameAndEmail));
      }
    }

    res = map.get("license");
    if(res != null) {
      license = (String) res;
    }
  }

  @SuppressWarnings("unchecked")
  private void retrievePayloads() throws IOException, RPCClient.MSFException {
    HashMap<String,ArrayList<String>> map = (HashMap<String,ArrayList<String>>) System.getMsfRpc().call("module.compatible_payloads", name);

    if(map == null)
      return;

    ArrayList<String> list = map.get("payloads");

    if(list == null)
      return;

    payloads.clear();

    for(String pName : list) {
      payloads.add(new Payload(pName));
    }

    if(payloads.size() == 1)
      currentPayload = payloads.get(0);
  }

  @SuppressWarnings("unchecked")
  public boolean launch() throws RPCClient.MSFException {
    if(System.getMsfRpc()==null || options.isEmpty()) {
      Logger.info("exploit launch failed: " + (options.isEmpty() ?
                      "missing options" : "RPC client not available"));
      return false;
    }

    try {
      Map<String,Object> opts = new HashMap<String, Object>();

      if(currentTarget != null)
        opts.put("TARGET", currentTarget);

      Collection<Option> allOptions = new ArrayList<>(options);

      if(currentPayload != null) {
        opts.put("PAYLOAD", currentPayload.getName());
        allOptions.addAll(currentPayload.getOptions());
      }

      for(Option o : allOptions) {
        opts.put(o.getName(), o.getValue());
      }

      Map<String, Object> res = (Map<String, Object>) System.getMsfRpc().call("module.execute","exploit", name, opts);

      if(res == null) {
        Logger.info("exploit launch failed: null response from RPC");
        return false;
      }

      if(res.get("job_id") == null) {
        Logger.info("exploit launch failed: 'job_id' not found or null");
        return false;
      }

      jobs.add((String) res.get("uuid"));

      return true;
    } catch (IOException e) {
      System.errorLogging(e);
    }
    return false;
  }

  public boolean tryLaunch() {
    try {
      return this.launch();
    } catch (Exception e) {
      return false;
    }
  }

  @Override
  public void setPort(Target.Port port) {
    super.setPort(port);

    for(Option o : options) {
      if(o.getName().equals("RPORT")) {
        o.setValue(Integer.toString(port.getNumber()));
        break;
      }
    }
  }

  @Override
  public void setParent(Target parent) {
    super.setParent(parent);

    for(Option o : options) {
      if(o.getName().equals("RHOST")) {
        o.setValue(parent.getAddress().getHostAddress());
        break;
      }
    }
  }

  public String getCurrentTarget() {
    return currentTarget;
  }

  public int getVersion() {
    return version;
  }

  @Override
  public String toString() {
    return name.substring(name.lastIndexOf('/')+1);
  }

  @Override
  public int getDrawableResourceId() {
    return R.drawable.exploit_msf;
  }

  public Collection<Option> getOptions() {
    return options;
  }

  public Collection<Payload> getPayloads() {
    return payloads;
  }

  public void setPayload(int index) {
    currentPayload = payloads.get(index);
  }

  public void setPayload(Payload payload) throws IllegalArgumentException {
    if (!payloads.contains(payload)) {
      throw new IllegalArgumentException(
              String.format("payload '%s' is not valid for exploit '%s'",
                      payload.getName(), name));
    }
    currentPayload = payload;
  }

  public Payload getCurrentPayload() {
    return currentPayload;
  }

  public Collection<String> getTargets() {
    return targets;
  }

  public void setTarget(int index) {
    currentTarget = targets.get(index);
  }

  public void setTarget(String target) throws IllegalArgumentException {
    if (!targets.contains(target)) {
      throw new IllegalArgumentException(
              String.format("target '%s' is not valid for exploit '%s'",
                      target, name));
    }
    currentTarget = target;
  }

  public Collection<String> getPlatforms() {
    return platforms;
  }

  public Collection<String> getJobs() {
    return jobs;
  }

  public Collection<Author> getAuthors() {
    return authors;
  }

  public Collection<String> getArchitectures() {
    return architectures;
  }

  public String getSourceUrl() {
    return "https://github.com/rapid7/metasploit-framework/blob/master/modules/" + name + ".rb";
  }

  public String getHistoryUrl() {
    return "https://github.com/rapid7/metasploit-framework/commits/master/modules/" + name + ".rb";
  }

  public String getLicense() {
    return license;
  }

  public Ranking getRank() {
    return rank;
  }

  public void copyTo(MsfExploit that) {
    that.architectures = architectures;
    that.authors = authors;
    that.currentPayload = currentPayload;
    that.currentTarget = currentTarget;
    that.description = description;
    that.id = id;
    that.jobs = jobs;
    that.license = license;
    that.name = name;
    that.options = options;
    that.payloads = payloads;
    that.version = version;
    that.platforms = platforms;
    that.rank = rank;
    that.summary = summary;
    that.synced = synced;
    that.targets = targets;
  }
}
